package com.fs.decorator;

import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.fs.Utils.GZIPUtils;
import com.fs.Utils.ResultUtil;
import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.util.Strings;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.List;


/**
 * 设置了响应体加密
 */
@Slf4j

@SuppressWarnings("all")
public class RewriteResponseDecorator extends ServerHttpResponseDecorator {
    public RewriteResponseDecorator(ServerHttpResponse delegate) {
        super(delegate);
    }

    @Override
    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
        log.info("" + this.getStatusCode().value());
        if (this.getStatusCode().equals(HttpStatus.OK) && body instanceof Flux) {
            //更新跨域
            this.getHeaders().setAccessControlAllowOrigin("*");
            this.getHeaders().setAccessControlAllowCredentials(true);
            Flux<? extends DataBuffer> fluxBody = Flux.from(body);
            return super.writeWith(fluxBody.buffer().map(dataBuffers ->
            {
                DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
                //零拷贝拼接
                DataBuffer join = dataBufferFactory.join(dataBuffers);
                byte[] oldContent = new byte[join.readableByteCount()];
                join.read(oldContent);
                //释放缓存
                DataBufferUtils.release(join);

                List<String> encodings = this.getHeaders().get("content-encoding");
                if (encodings != null) {
                    for (String encoding : encodings) {
                        if (encoding.equals("gzip")) {
                            //GZIP解压
                            oldContent = GZIPUtils.uncompress(oldContent);
                            //重写Body后用GZIP压缩
                            byte[] newOcontent = encode(oldContent);
                            log.debug(new String(oldContent, StandardCharsets.UTF_8));
                            newOcontent = GZIPUtils.compressByByte(newOcontent);
                            //更新内容长度
                            this.getHeaders().setContentLength(newOcontent.length);
                            return dataBufferFactory.wrap(newOcontent);
                        }
                    }
                    ResultUtil resultUtil = new ResultUtil();
                    resultUtil.setStatus(402);
                    resultUtil.setMessage("ERROR!");
                    resultUtil.setData("no content-encoding type!");
                    return dataBufferFactory.wrap(JSONObject.toJSONBytes(resultUtil, SerializerFeature.WriteSlashAsSpecial, SerializerFeature.PrettyFormat));
                } else {
                    ResultUtil resultUtil = new ResultUtil();
                    resultUtil.setStatus(0);
                    resultUtil.setMessage("OK");
                    resultUtil.setData(Strings.fromUTF8ByteArray(oldContent));
//                    DefaultDataBuffer wrap = new DefaultDataBufferFactory().wrap(JSONObject.toJSONBytes(resultUtil));
                    DataBuffer wrap = dataBufferFactory.wrap(JSONObject.toJSONBytes(resultUtil, SerializerFeature.WriteSlashAsSpecial, SerializerFeature.PrettyFormat));
                    //设置响应体长度，
                    getDelegate().getHeaders().setContentLength(wrap.capacity());
                    return wrap;
                }
            }));
        }
        return super.writeWith(body);
    }

    byte[] encode(byte[] content) {
        ResultUtil resultUtil = new ResultUtil();
        resultUtil.setStatus(0);
        resultUtil.setMessage("OK");
        resultUtil.setData(Strings.fromUTF8ByteArray(content));
        String res = JSONObject.toJSONString(resultUtil, SerializerFeature.WriteSlashAsSpecial, SerializerFeature.PrettyFormat);
        byte[] resultUtilContent = Strings.toUTF8ByteArray(res);
        return resultUtilContent;
    }

    @Override
    public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
        return writeWith(Flux.from(body).flatMapSequential(p -> p));
    }
}
